#
# Created: 2025-08-01

import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as st
import pandas as pd
import pingouin as pg

# 设置全局字体
plt.rcParams['font.family'] = 'Times New Roman'  # 默认使用Times New Roman

# Step 1: 读取本地数据
def read_synthetic_data():
    df=pd.read_csv("synthetic_data.csv")
    categories = ["Type A", "Type B", "Type C", "Type D"]
    y_data = [df[df["Group"] == cat]["Value"].values for cat in categories]

    jitter = 0.04  # x轴抖动的尺度
    x_jittered = [
        np.full(len(y), i) + st.t(df=6, scale=jitter).rvs(len(y))  # 每组加抖动
        for i, y in enumerate(y_data)
    ]

    return categories, y_data, x_jittered

# Step 2: 设定绘图风格与颜色
def setup_visuals():
    colors = {
        "BG_WHITE": "white",           # 整体背景色
        "GREY_LIGHT": "black",         # 黑色（坐标轴等）
        "GREY50": "#999999",           # 中灰色（辅助线）
        "BLUE_DARK": "#1B2838",        # 深蓝色（标题）
        "BLUE": "#2a475e",             # 蓝色（副标题）
        "BLACK": "#282724",            # 黑色（轮廓线）
        "GREY_DARK": "#747473",        # 深灰（箱线图元素）
        "RED_DARK": "#850e00"          # 红色（均值点）
    }
    color_scale = ["#1B9E77", "#D95F02", "#7570B3", "#E7298A"]  # 每组数据点颜色
    positions = [0, 1, 2, 3]  # x轴位置
    hlines = [40, 50, 60, 70]  # 背景参考线y坐标
    return colors, color_scale, positions, hlines

# Step 3: 添加每组均值点与文字标注
def add_statistical_annotations(ax, y_data, colors):
    for i, y in enumerate(y_data):
        mean = np.mean(y)  # 计算当前组的均值

        # ① 添加均值点（红色圆点）
        ax.scatter(i, mean,
                   s=250,  # 点的大小
                   color=colors["RED_DARK"],  # 使用颜色字典中预定义的深红色
                   zorder=3)  # 图层高度，确保点显示在其他图层之上

        # ② 添加从均值点向右延伸的短虚线（dashdot）
        ax.plot([i, i + 0.25],  # 横坐标从当前组位置到右边0.25
                [mean, mean],  # 均值水平线
                ls="dashdot",  # 虚线样式为 dash-dot
                color="black",  # 线的颜色
                zorder=3)  # 图层高度

        # ③ 添加均值的文字标注（μ = xx.xx），使用 latex 公式格式
        ax.text(i + 0.25, mean,  # 坐标在均值右侧
                f"$\\hat{{\\mu}}$ = {mean:.2f}",  # 显示均值到小数点后两位
                fontsize=13,  # 字体大小
                va="center",  # 垂直对齐方式：居中
                bbox=dict(  # 设置文字的边框样式
                    facecolor="white",  # 背景白色
                    edgecolor="black",  # 边框黑色
                    boxstyle="round",  # 圆角矩形
                    pad=0.15),  # 内边距
                zorder=10)  # 保证文字浮在最上层

# Step 4: 添加 Games-Howell 显著性p值文字（非标注线）
def games_howell_annotations(ax, y_data, group_labels, positions, buffer=2.0, step=2.0):
    """
    使用 Games-Howell 方法计算两两组之间的显著性差异，并在图上方绘制对应的 p 值文字和连接线。
    参数：
    ax           -- matplotlib 的坐标轴对象
    y_data       -- 每组的 y 值数据（list of arrays）
    group_labels -- 每组的标签，如 ['Type A', 'Type B', 'Type C', 'Type D']
    positions    -- 每组在x轴上的位置，通常为 [0, 1, 2, 3]
    buffer       -- 所有数据最大值向上偏移量（避免被点遮住）
    step         -- 每组之间纵向间距，防止标注重叠
    """
    # ① 将所有数据合并成一列
    data = np.concatenate(y_data)

    # ② 构造对应的组别标签列，与数据一一对应
    groups = np.concatenate([
        [group_labels[i]] * len(y)  # 重复每个标签 len(y) 次
        for i, y in enumerate(y_data)
    ])

    # ③ 构建用于统计检验的 DataFrame，列名为 'data' 和 'group'
    df = pd.DataFrame({'data': data, 'group': groups})

    # ④ 执行 Games-Howell 多重比较检验（适用于方差不齐、样本不等）
    result = pg.pairwise_gameshowell(data=df, dv='data', between='group')

    # ⑤ 按照 p 值从小到大排序，便于逐层向上绘制
    result = result.sort_values(by="pval")

    # ⑥ 获取所有 y 值的最大值，并上移 buffer 单位作为起始绘图位置
    y_base = max([np.max(y) for y in y_data]) + buffer

    # ⑦ 逐行绘制每一组显著性比较的连线和文字
    for idx, row in result.iterrows():
        g1, g2 = row['A'], row['B']  # 获取比较的组名
        pval = row['pval']  # 获取对应的 p 值

        x1 = positions[group_labels.index(g1)]  # 找到组A的 x 坐标
        x2 = positions[group_labels.index(g2)]  # 找到组B的 x 坐标
        y = y_base + idx * step  # 当前 p 值标注线的 y 高度（避免重叠）

        # ⑧ 绘制连接线（先竖直上升，再水平连接，再竖直下降）
        ax.plot([x1, x1, x2, x2],  # 连接线横坐标
                [y - 0.2, y, y, y - 0.2],  # 对应纵坐标，形成“∩”结构
                c="black")  # 线条颜色为黑色

        # ⑨ 在两组中间位置添加 p 值文字说明，格式为科学计数法
        ax.text((x1 + x2) / 2, y + 0.2,  # 文字位置：横坐标中点、纵坐标略高
                f"$p_{{\\rm{{Holm-corrected}}}}$ = {pval:.2e}",  # LaTeX 格式显示
                ha="center", va="bottom", fontsize=11)

# Step 5: 美化图像整体布局与字体
def customize_layout(ax, categories, y_data, colors, hlines, positions):
    # 去掉顶部和右侧边框
    ax.spines["right"].set_color("none")
    ax.spines["top"].set_color("none")

    # 设置左下边框颜色和线宽
    ax.spines["left"].set_color(colors["GREY_LIGHT"])
    ax.spines["bottom"].set_color(colors["GREY_LIGHT"])
    ax.spines["left"].set_linewidth(2)
    ax.spines["bottom"].set_linewidth(2)

    # 主刻度设置
    ax.set_yticks(hlines)  # 如 [40, 50, 60, 70]
    ax.set_yticklabels([f"{y} mm" for y in hlines], size=15)
    ax.set_ylabel("Length (mm)", size=18, weight="bold")

    xlabels = [f"{cat}\n(n={len(y_data[i])})" for i, cat in enumerate(categories)]
    ax.set_xticks(positions)
    ax.set_xticklabels(xlabels, size=15)
    ax.set_xlabel("Synthetic Category", size=18, weight="bold")

    # 设置刻度方向朝外，线长和粗细
    ax.tick_params(axis="both", direction="out", length=6, width=2)

    # 设置次刻度（minor ticks）
    from matplotlib.ticker import AutoMinorLocator
    ax.yaxis.set_minor_locator(AutoMinorLocator(n=2))  # 每个主刻度间插入1个次刻度
    ax.tick_params(which='minor', length=4, width=1, color='black')


# Step 6: 添加标题与说明文字
def add_titles_and_captions(fig, ax, colors):
    fig.suptitle("Distribution of Length across Synthetic Categories",ha="center", fontsize=24,
                 fontname="Lobster Two", color=colors["BLUE"], weight="bold")
    ax.set_title("Synthetic F-test: F=221.3, p=1.2e-42, ω²=0.82",
                 loc="left", ha="left", fontsize=18, color=colors["BLUE_DARK"])
    fig.text(0.7, 0.02, "Note: Simulated data; comparisons are illustrative only.", fontsize=10)

# Step 7: 创建图表主函数
def create_plot(categories, y_data, x_jittered, colors, color_scale, positions, hlines):
    fig, ax = plt.subplots(figsize=(12, 8))
    fig.patch.set_facecolor(colors["BG_WHITE"])
    ax.set_facecolor(colors["BG_WHITE"])

    # 背景横线
    for h in hlines:
        ax.axhline(h, color=colors["GREY50"], ls=(0, (5, 5)), alpha=0.8)

    # 小提琴图（轮廓）
    violins = ax.violinplot(y_data, positions=positions, widths=0.45,
                            bw_method="silverman", showmeans=False,
                            showmedians=False, showextrema=False)
    for pc in violins["bodies"]:
        pc.set_facecolor("none")  # 设置为透明
        pc.set_edgecolor(colors["BLACK"])  # 外框颜色
        pc.set_linewidth(1.5)

    # 箱线图（中位数、箱体、须）
    ax.boxplot(y_data, positions=positions, showfliers=False,
               showcaps=False,
               medianprops=dict(linewidth=3, color=colors["GREY_DARK"]),
               whiskerprops=dict(linewidth=1.5, color=colors["GREY_DARK"]),
               boxprops=dict(linewidth=1.5, color=colors["GREY_DARK"]))

    # 抖动散点
    for x, y, c in zip(x_jittered, y_data, color_scale):
        ax.scatter(x, y, s=100, color=c, alpha=0.4)

    add_statistical_annotations(ax, y_data, colors)
    games_howell_annotations(ax, y_data, categories, positions)
    customize_layout(ax, categories, y_data, colors, hlines, positions)
    add_titles_and_captions(fig, ax, colors)

    # 保存图像（高分辨率）
    fig.savefig("synthetic_violin_plot.png", dpi=300)
    fig.savefig("synthetic_violin_plot.pdf")
    print("✅ 图像保存为 synthetic_violin_plot.png 和 synthetic_violin_plot.pdf")

    return fig, ax

# Step 8: 主函数调用
if __name__ == "__main__":
    categories, y_data, x_jittered = read_synthetic_data()
    colors, color_scale, positions, hlines = setup_visuals()
    fig, ax = create_plot(categories, y_data, x_jittered, colors, color_scale, positions, hlines)
    plt.show()